from warnings import catch_warnings
import torch
from torch.nn.modules import conv
from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn.modules.instancenorm import InstanceNorm2d
import torchvision
from torch import nn

def conv3X3(in_channels, out_channels, stride = 1):
    return nn.Conv2d(in_channels, out_channels, stride = 1, padding = 1, kernel_size = 3, bias = False)


class residual(nn.Module):
    def __init__(self, channels):
        super(residual, self).__init__()
        self.res = nn.Sequential(
            conv3X3(channels, channels),
            # nn.BatchNorm2d(channels),
            nn.ReLU(True),
            conv3X3(channels, channels),
            # nn.BatchNorm2d(channels)
        )
        self.RELU = nn.ReLU(True)
    
    def forward(self, x):
        res = x
        output = self.res(x)
        output += res
        output = self.RELU(output)
        return output

class deconv(nn.Module):
    def __init__(self, in_channels, out_channels, stride = 2):
        super(deconv, self).__init__()
        self.deconvBlock = nn.Sequential(
            nn.ConvTranspose2d(in_channels = in_channels, out_channels = out_channels, stride = stride, kernel_size = 4, padding = 1, bias = False),
            nn.InstanceNorm2d(out_channels),
            nn.ReLU(True)
        )
    
    def forward(self, x):
        output = self.deconvBlock(x)
        return output


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.downSample = nn.Sequential(
            nn.Conv2d(3, 32, 9, 2, 4, bias = False),
            nn.InstanceNorm2d(32),
            nn.ReLU(True),
            nn.Conv2d(32, 64, 3, 2, 1, bias = False),
            nn.InstanceNorm2d(64),
            nn.ReLU(True),
            nn.Conv2d(64, 128, 3, 2, 1, bias  = False),
            nn.InstanceNorm2d(128),
            nn.ReLU(True)
        )

        self.residualBlocks = nn.Sequential(
            residual(128),
            residual(128),
            residual(128),
            residual(128),
            residual(128)
        )

        self.deconvBlocks = nn.Sequential(
            deconv(128, 64),
            deconv(64, 32)
        )

        self.lastDeconvBlock = nn.Sequential(
            nn.ConvTranspose2d(in_channels = 32, out_channels = 3, stride = 2, kernel_size = 4, padding = 1, bias = False),
            nn.InstanceNorm2d(3),
            nn.Tanh()
        )
    
    def forward(self, x):
        afterDownSampling = self.downSample(x)
        afterRes = self.residualBlocks(afterDownSampling)
        afterDeconvs = self.deconvBlocks(afterRes)
        y = self.lastDeconvBlock(afterDeconvs)
        return (y + 1) / 2